"""
Right-click Menu Manager Script

This Python script allows users to manage Windows right-click context  menu items. It can be used to create or delete
custom menu items in the context (right-click) menu that appears when you right-click on a file or folder in the
Windows Explorer.

The script works by modifying the Windows registry, adding or removing entries under the following registry path:
    HKEY_CLASSES_ROOT\\*\\shell\\{menu_name}

The menu items created by this script will script trigger the execution of a specified executable the file (for example,
a program or script) when selected from the right-click memu.

Usage:

1. Create a Right-Click Menu Item:
   To create a new menu item, run the script with the 'create' action, provide the name for the menu item. and specify
the path to the executable that should run when the menu item is clicked.

   Example:
   python right_click_menu_manager.py create MyMenuItem --exe_path "C:\\Path\\To\\YourApp.exe"

   This will create a new menu item called 'MyMenuItem' that runs 'YourApp.exe' when selected

2. Delete a Right-Click Menu Item:
   To delete an existing menu item, run the script with the 'delete' action and provide the name of the menu item you
want to remove.

   Example:
   pytho right_click_menu_manager.py delete MyMenuItem

   This will delete the 'MyMenuItem' from the right-click context menu.

Arguments:
- action (str): The action to perform, either 'create' or 'delete'.
- menu_name (str): The name of the right-click menu item to create or delete.
- --exe_path (str, optional): The path to the executable to associate with the menu item. This is required only when
using the 'create' action

Note:
- The script modifies Windows registry keys to create or remove right-click menu items. Ensure you have appropriate
permissions (administrator rights may be required).
- Use the 'create' action to add a custom right-click item that runs a program, and the 'delete' action to remove it.
- The script is cross-platform, but the functionality is tailored for Windows systems.

Example use cases:
- Quickly add a right-click option to open a specific file with a custom application.
- Remove unnecessary or unwanted right-click context menu items.

"""
import argparse
import winreg as reg
import os
import concurrent.futures


def build_right_click_menu_with_threaded(func, pos, arg):
    with concurrent.futures.ThreadPoolExecutor() as executor:
        executor.map(lambda po: func(po, arg), pos)


def recursion_delete(sub_key, key):
    query_key = reg.OpenKey(key, sub_key)
    key_size, _, _ = reg.QueryInfoKey(query_key)
    for i in range(key_size):
        name = reg.EnumKey(query_key, i)
        recursion_delete(sub_key + '\\' + name, key)
    reg.DeleteKey(key, sub_key)


def add_right_click_menu_with_pos(pos, exe_path):
    # noinspection PyBroadException
    try:
        command_key = pos + '\\command'
        reg.CreateKey(reg.HKEY_CLASSES_ROOT, pos)
        reg.CreateKey(reg.HKEY_CLASSES_ROOT, command_key)
        reg.SetValue(reg.HKEY_CLASSES_ROOT, pos, reg.REG_SZ, '执行 {0}'.format(pos.split('\\')[-1]))
        reg.SetValue(reg.HKEY_CLASSES_ROOT, command_key, reg.REG_SZ, f'"{exe_path}" "%1"')

    except Exception as e:
        raise Exception(f'创建 {pos} 右键菜单项失败: {e}')


def add_right_click_menu(position, menu_name, exe_path):
    build_right_click_menu_with_threaded(add_right_click_menu_with_pos,
                                         [f'{pos}\\{"Shell" if pos == "DesktopBackground" else "shell"}\\{menu_name}'
                                          for pos
                                          in position], exe_path)


def remove_right_click_menu(position, menu_name):
    """
    DeleteKey(Warning): This method can not delete keys with sub keys.
    """
    build_right_click_menu_with_threaded(recursion_delete,
                                         [f'{pos}\\{"Shell" if pos == "DesktopBackground" else "shell"}\\{menu_name}'
                                          for pos
                                          in position], reg.HKEY_CLASSES_ROOT)


def main():
    parser = argparse.ArgumentParser(description='管理右键菜单项')
    parser.add_argument('position', choices=['DesktopBackground', 'Directory', '*'], nargs='+',
                        help='选择一个或多个位置: black, dir, file')
    parser.add_argument('action', choices=['create', 'delete'], help='操作类型: create 或 delete')
    parser.add_argument('menu_name', help='右键菜单项名称')
    parser.add_argument('--exe_path', help='执行的程序路径 (仅在创建时使用)', default='C:\\Path\\To\\YourApp.exe')

    args = parser.parse_args()

    if args.action == 'create':
        if os.path.exists(args.exe_path):
            add_right_click_menu(args.position, args.menu_name, args.exe_path)
        else:
            raise FileNotFoundError(f'执行的程序路径不存在...')
    elif args.action == 'delete':
        remove_right_click_menu(args.position, args.menu_name)


if __name__ == '__main__':
    # execute admin
    main()
